#!/bin/bash

# import Global Variable
. /usr/local/vnc/conf/config

echo_err() {
  echo "$*" >&2
}

print_vnc_conf(){

  echo "---- VNC Configuration ----"
  echo VNC_ROOT=${VNC_ROOT}
  echo VNC_SESSION_DIR=${VNC_SESSION_DIR}
  echo VNC_CONF_DIR=${VNC_CONF_DIR}
  echo VNC_CONF_DEFAULT=${VNC_CONF_DEFAULT}
  echo VNC_HOSTNAME=${VNC_HOSTNAME}
  echo VNC_SERVER=${VNC_SERVER}
  echo VNC_PASSWD=${VNC_PASSWD}
  echo VNC_CLASS_DIR=${VNC_CLASS_DIR}
  echo VNC_CODE=${VNC_CODE}
  echo VNC_ARCHIVE=${VNC_ARCHIVE}
  echo VNC_LOCALE=${VNC_LOCALE}
  echo

}

set_lang(){

  # Setting locale
  if [ -z "${VNC_LOCALE}" ] ; then
    VNC_LOCALE="C"
  fi
  LANG="${VNC_LOCALE}"; export LANG
  LANGUAGE="${VNC_LOCALE}"; export LANGUAGE
  LC_ALL="${VNC_LOCALE}"; export LC_ALL
  

}

check_var(){

  # Check VNC_ROOT
  if [ -z "${VNC_ROOT}" ]; then
    echo_err "VNC_ROOT not set"
    return 1
  fi
  if [ ! -d ${VNC_ROOT} ]; then
    echo_err "VNC_ROOT(${VNC_ROOT}) not exist"
    return 1
  fi

  # Check VNC_SESSION_DIR
  if [ -z "${VNC_SESSION_DIR}" ]; then
    echo_err "VNC_SESSION_DIR not set"
    return 1
  fi
  if [ ! -d ${VNC_SESSION_DIR} ]; then
    echo_err "VNC_SESSION_DIR(${VNC_SESSION_DIR}) not exist"
    return 1
  fi

  # Check VNC_CONF_DIR
  if [ -z "${VNC_CONF_DIR}" ]; then
    echo_err "VNC_CONF_DIR not set"
    return 1
  fi
  if [ ! -d ${VNC_CONF_DIR} ]; then
    echo_err "VNC_CONF_DIR(${VNC_CONF_DIR}) not exist"
    return 1
  fi

  # Check VNC_CONF_DEFAULT
  if [ -z "${VNC_CONF_DEFAULT}" ]; then
    echo_err "VNC_CONF_DEFAULT not set"
    return 1
  fi
  if [ ! -f ${VNC_CONF_DEFAULT} ]; then
    echo_err "VNC_CONF_DEFAULT(${VNC_CONF_DEFAULT}) not exist"
    return 1
  fi

  # Check VNC_HOSTNAME
  if [ -z "${VNC_HOSTNAME}" ]; then
    echo_err "VNC_HOSTNAME not set"
    return 1
  fi
  if ! ping -c 1 ${VNC_HOSTNAME} > /dev/null 2>&1; then
    echo_err "VNC_HOSTNAME(${VNC_HOSTNAME}) not connected"
    return 1
  fi

  # Check VNC_SERVER
  if [ -z "Z${VNC_SERVER}" ]; then
    echo_err "VNC_SERVER not set"
    return 1
  fi
  if [ ! -x ${VNC_SERVER} ]; then
    echo_err "VNC_SERVER(${VNC_SERVER}) not executable"
    return 1
  fi

  # Check VNC_PASSWD
  if [ -z "${VNC_PASSWD}" ]; then
    echo_err "VNC_PASSWD not set"
    return 1
  fi
  if [ ! -x ${VNC_PASSWD} ]; then
    echo_err "VNC_PASSWD(${VNC_PASSWD}) not executable"
    return 1
  fi

  # Check VNC_CLASS_DIR
  if [ -z "${VNC_CLASS_DIR}" ]; then
    echo_err "VNC_CLASS_DIR not set"
    return 1
  fi
  if [ ! -d ${VNC_CLASS_DIR} ]; then
    echo_err "VNC_CLASS_DIR(${VNC_CLASS_DIR}) not exist"
    return 1
  fi

  # Check VNC_CODE
  if [ -z "${VNC_CODE}" ]; then
    echo_err "VNC_CODE not set"
    return 1
  fi
  # Commont out VNC_CODE
  # Because in some kind of Linux, VncViewer.class exist in VNC_ARCHIVE (jar package)
  #
  #if [ ! -f ${VNC_CLASS_DIR}/${VNC_CODE} ]; then
  #  echo_err "VNC_CODE(${VNC_CODE}) not exist"
  #  return 1
  #fi

  # Check VNC_ARCHIVE
  if [ -z "${VNC_ARCHIVE}" ]; then
    echo_err "VNC_ARCHIVE not set"
    return 1
  fi
  if [ ! -f ${VNC_CLASS_DIR}/${VNC_ARCHIVE} ]; then
    echo_err "VNC_ARCHIVE(${VNC_ARCHIVE}) not exist"
    return 1
  fi

  # Check VNC_LOCALE
  if [ -z "${VNC_LOCALE}" ]; then
    VNC_LOCALE="C"
  fi
  LANG="${VNC_LOCALE}"; export LANG
  LANGUAGE="${VNC_LOCALE}"; export LANGUAGE
  LC_ALL="${VNC_LOCALE}"; export LC_ALL

  return 0
}

check_job(){
  source /etc/profile.d/dawning.sh
  local squeuecmd="squeue"
  if [ -n "$SLURM_HOME" ];then
        squeuecmd="$SLURM_HOME/bin/squeue"
  fi
  local job=`$squeuecmd --job $1 -h 2>&1`
  if [ -z "$job" ];then
        return 1
  else
		local invalidjob=`echo "${job}"|grep "Invalid job id specified"`
		if [ ! -z "$invalidjob" ];then
			return 1
		else
			return 0
		fi
  fi
}

check_jobcheduler(){
  source /etc/profile.d/dawning.sh
  local sinfocmd="sinfo"
  if [ -n "$SLURM_HOME" ];then
        sinfocmd="$SLURM_HOME/bin/sinfo"
  fi
  timeout 60 $sinfocmd > /dev/null 2>&1
}

#
# for random string
#
get_random_str(){

  local randomstr=""

  # Length of random str
  local strlen=8

  #############
  # Method Zero:
  #             : limit-1: static string
  #############
  #randomstr="dawning123"

  #############
  # Method One:
  #             : limit-1: length maybe difference
  #############
  #randomstr="$RANDOM"

  #############
  # Method Two:
  #             : limit-1: depend on utils "od" (coreutils)
  #             : limit-2: only 8-bit string supported
  #############
  #randomstr=$(od -N 4 -t x4 /dev/random | head -1 | awk '{print $2}')
  # Change random to urandom, because random may take too long time before return (at CentOS7)
  randomstr=$(od -N 4 -t x4 /dev/urandom | head -1 | awk '{print $2}')

  #############
  # Method Three:
  #             : limit-1: depend on utils "md5sum" (coreutils)
  #             : limit-2: strlen should less than 33 (md5sum)
  #############
  #randomstr=$(cat /dev/urandom | head -1 | md5sum | head -c $strlen)

  #############
  # Method Four:
  #             : limit-1: depend on utils "strings" (binutils)
  #             : limit-2: response time is longer than other method
  #############
  #randomstr=$(cat /dev/urandom | sed 's/[^a-zA-Z0-9]//g' | strings -n $strlen | head -c $strlen)

  echo $randomstr
  
}

#
# Generate random passwd filename
#
get_pass_file(){

  #
  # final passfile name
  #
  local passfile=""

  #
  # random string, for passfile name
  #
  local ramstr=`get_random_str`

  #
  # base name for passwd file
  #
  local basefile="${HOME}/.vnc/passwd"

  #
  # get filename, format {base file}.{user}.{random}
  #
  #passfile=${basefile}.`whoami`.`date '+%Y%m%d%H%M%S'`.${ramstr}
  passfile=${basefile}.`date '+%Y%m%d%H%M%S'`.${ramstr}
  echo $passfile
  
}

#
# Generate random string for vnc-passwd
#
get_pass(){

  local pass=`get_random_str`

  echo $pass
  
}

#
# query proc pid of specified session
# usage: pid_of_user_session {user name} {session id}
#
pid_of_user_session(){

  if [ $# -ne 2 ]; then
    echo_err "invalid arguments"
    return
  fi
  local uname=$1
  local sid=$2
  if ! id $uname >&/dev/null; then
    echo_err "user $uname not exist"
    #return
  fi
  if
	  lsb_release -d|grep Linx >/dev/null 2>&1
  then
	#Linx
  	ps -eo pid,user:32,args | grep -w "Xvnc4 :${sid}" | grep -v 'grep' | awk -v USER=${uname} '{if ($2 == USER) print $1}'
  else
  	ps -eo pid,user:32,args | grep -w "Xvnc :${sid}" | grep -v 'grep' | awk -v USER=${uname} '{if ($2 == USER) print $1}'
  fi
}

#
# clear warnning session
# usage: clear_warning {vnc start log file}
# for example: /tmp/xxx
#
clear_warning(){

  local idlist="`awk -F: 'NF==2{if($1 ~/Remove this file if there is no X server/){print $2}}' $1`"
  #echo $idlist
  if [ ! -z "${idlist}" ]; then
    for i in ${idlist};
    do
      lock1=/tmp/.X11-unix/X$i
      lock2=/tmp/.X$i-lock
      # care socket file
      if [ -S ${lock1} ]; then
	rm -f ${lock1}
      fi
      # care socket file
      if [ -f ${lock2} ]; then
	rm -f ${lock2}
      fi
    done
  fi

}

vnc_start(){
  UNAME="`whoami`"
  VNC_PROXYSERVER=`getvnchost`
  #VNC_PROXYSERVER="gv244"
  if [ $? -eq 111 ]; then
    exit 111
  fi
  job_id="$1"
  queue="$2"
  app_name="$3"
  _result=`ssh -o StrictHostKeyChecking=no -o BatchMode=yes ${VNC_PROXYSERVER} "bash /usr/local/vnc/bin/vnc_create '${UNAME}' '$job_id' '' '$queue' '$app_name'"`
  echo "${_result}"
  #/usr/local/vnc/bin/managercreate ${UNAME} $1
  vncresult="sessionInfo=${_result}"
  curl --connect-timeout 5 -m 20 -X POST -d "${vncresult}" ${MANAGER_NODE_URI}/desktopandapp/api/vncs/vncdbinfo &>/dev/null
  return $?
}

vnc_stop(){
  UNAME="`whoami`"
  #ssh -o StrictHostKeyChecking=no -o BatchMode=yes $2 /usr/local/vnc/bin/vnc_delete $UNAME $1 > /dev/null 2>&1
  ssh -o StrictHostKeyChecking=no -o BatchMode=yes $2 /usr/local/vnc/bin/vnc_delete $UNAME $1
  curl --connect-timeout 5 -m 20 -X DELETE  ${MANAGER_NODE_URI}/desktopandapp/api/vncs/vncdbinfo?sid=$1\&hostname=$2 &>/dev/null
  return $?
}

getvnchost(){
_executename=''
_min=0
_index=1
_mem=0
_fix=$1
#
# Check all VNC_PROXYS servers
#

# Query "VNC_PROXYS" configuration from $VNC_HOSTNAME
if [ -z "$_fix" ];then
PROXYS_AND_NAME=$(ssh -o StrictHostKeyChecking=no -o BatchMode=yes ${VNC_HOSTNAME} '/usr/local/vnc/bin/vnc_config | grep -w VNC_PROXYS ; cat /tmp/vnc_name')
else
PROXYS_AND_NAME=$(/usr/local/vnc/bin/vnc_config | grep -w VNC_FIX_PROXYS ; cat /tmp/vnc_fixname)
fi
if [ ! -z "$PROXYS_AND_NAME" ]; then
  if [ -z "$_fix" ];then
  VNC_PROXYS=`echo "${PROXYS_AND_NAME}" |grep -w VNC_PROXYS | awk -F= 'NF==2{print $2}'`
  else
  VNC_PROXYS=`echo "${PROXYS_AND_NAME}" |grep -w VNC_FIX_PROXYS | awk -F= 'NF==2{print $2}'`
  fi
  VNC_CURRENT_PROXY_NAME=`echo "${PROXYS_AND_NAME}"|grep -w CURRENT_NAME | awk -F= 'NF==2{print $2}'`
  if [ -z "${VNC_CURRENT_PROXY_NAME}" ]; then
    VNC_CURRENT_PROXY_NAME=''
  elif [[ "${VNC_PROXYS}" =~ "${VNC_CURRENT_PROXY_NAME}" ]];then
    :
  else
    VNC_CURRENT_PROXY_NAME=''
  fi
fi
if [ -n "${VNC_CURRENT_PROXY_NAME}" ];then
	var1=$(echo ${VNC_PROXYS} | awk -F"${VNC_CURRENT_PROXY_NAME} " '{print $1}')
	var2=$(echo ${VNC_PROXYS} | awk -F"${VNC_CURRENT_PROXY_NAME} " '{print $2}')
    VNC_PROXYS=$(echo "${var2} ${var1} ${VNC_CURRENT_PROXY_NAME}")
fi
for i in ${VNC_PROXYS};
do
    _date="`date +'%Y-%m-%d%H:%M:%S'`"
    UNAME="`whoami`"
    vnc_load=`ssh -o ConnectTimeout=5 -o StrictHostKeyChecking=no -o BatchMode=yes $i /usr/local/vnc/util/getcpuandmem.sh -q`
    _cpu=`echo ${vnc_load} |grep "vnc_load" | awk '{print $2}'`
    _mem=`echo ${vnc_load} | grep "vnc_load" | awk '{print $3}'`
    if [ "${_cpu}" == "" ]; then
      _cpu=100
    fi
    if [ "${_mem}" == "" ]; then
      _mem=100
    fi
    #set the first vnc proxy as the min
    if [ "${_index}" -eq 1 ];then
      _executename=$i
      _min_cpu=${_cpu}
      _min_mem=${_mem}
      _index=2
    fi

    #sleep 1

    #compare current cpu usage with min and make sure it is lower than 90
    _result=`echo "${_min_cpu}>${_cpu}"|bc`
    _result2=`echo "${_cpu}<${VNC_LIMIT}"|bc`
    if [ "${_result}" -eq 1 -a "${_result2}" -eq 1 ]; then
      _compare=`echo "${_mem}<${VNC_LIMIT}"|bc`
      if [ "${_compare}" -eq 1 ]; then
        _min_cpu=${_cpu}
        _min_mem=${_mem}
        _executename=$i
      fi
    fi

    #compare min cpu usage with min_limit, if lower, just use this one
    compare_mincpu=`echo "${_min_cpu}<${VNC_MIN_LIMIT}"|bc`
    compare_maxmem=`echo "${_min_mem}<${VNC_LIMIT}"|bc`
    if [ "${compare_mincpu}" -eq 1 -a "${compare_maxmem}" -eq 1 ];then
      _executename=$i
	  VNC_CURRENT_PROXY_NAME=${_executename}
      break
    fi
done
if [ -z "$_fix" ];then
ssh -o StrictHostKeyChecking=no -o BatchMode=yes ${VNC_HOSTNAME} 'touch /tmp/vnc_name >/dev/null 2>&1;chmod 777 /tmp/vnc_name >/dev/null 2>&1;echo 'CURRENT_NAME=${VNC_CURRENT_PROXY_NAME}' > /tmp/vnc_name'
else
touch /tmp/vnc_fixname >/dev/null 2>&1;chmod 777 /tmp/vnc_fixname >/dev/null;echo "CURRENT_NAME=${VNC_CURRENT_PROXY_NAME}" > /tmp/vnc_fixname
fi
cpu_compare=`echo "${_min_cpu}<${VNC_LIMIT}"|bc`
mem_compare=`echo "${_min_mem}<${VNC_LIMIT}"|bc`
if [[ "${cpu_compare}" == 1 && "${mem_compare}" == 1 ]]; then
  echo ${_executename};
else
  echo_err "all vnc proxys are overloaded or ssh failed"
  exit 111
fi
}

gethost(){
_executename=''
_min=0
_index=1

#
# Check all VNC_PROXYS servers
#

# Method One: not worked
#VNC_PROXYS=`getVncProxyServers`

# Method Two: not worked
#VNC_PROXYS already setted in /usr/local/vnc/conf/config

# Method Three: worked
# Query "VNC_PROXYS" configuration from $VNC_HOSTNAME
ALL_PROXYS=$(ssh -o StrictHostKeyChecking=no -o BatchMode=yes ${VNC_HOSTNAME} /usr/local/vnc/bin/vnc_config | grep -w VNC_PROXYS | awk -F= 'NF==2{print $2}')
if [ ! -z "$ALL_PROXYS" ]; then
  VNC_PROXYS=$ALL_PROXYS
fi

for i in ${VNC_PROXYS};
do
_date="`date +'%Y-%m-%d%H:%M:%S'`"
#echo "${_date}"
#echo "${arg_user}"
UNAME="`whoami`"
#echo $i
ssh -o StrictHostKeyChecking=no -o BatchMode=yes $i /usr/local/vnc/bin/vnc_query root >> /tmp/$i."${_date}".txt
_num=`cat /tmp/$i.${_date}.txt | wc -l`
#echo ${_date}
if [ "${_index}" -eq 1 ];then
_executename=$i
_min="${_num}"
_index=2
fi

sleep 1
#echo ${_num}
if [ "${_num}" -lt "${_min}" ]; then
  _min=${_num}
_executename=$i
#echo ${_mun}
fi

done
echo ${_executename};
}


function clean_passwdfile_and_xstartup()
{

  #
  # Griview3.2.0, add _VNC_PASSWDFILE and _VNC_XSTARTUP
  #
  
  # delete _VNC_PASSWDFILE (which is not the default ${HOME}/.vnc/passwd)
  # worked at Gridview3.2.0 and later
  #
  if [ ! -z "${_VNC_PASSWDFILE}" ] && [ "${_VNC_PASSWDFILE}" != "${HOME}/.vnc/passwd" ]; then
    rm -rf "${_VNC_PASSWDFILE}" > /dev/null 2>&1
  fi
  
  #
  # delete _VNC_XSTARTUP (which is not the default ${HOME}/.vnc/xstartup)
  # NOT worked at Gridview3.2.0, because ${HOME}/.vnc/xstartup is used
  #
  if [ ! -z "${_VNC_XSTARTUP}" ] && [ "${_VNC_XSTARTUP}" != "${XSTARTUP_USER}" ]; then
    rm -rf "${_VNC_XSTARTUP}" > /dev/null 2>&1
  fi

}

function touch_pass_file()
{

  local findname=0
  local filename=""
  for i in `seq 1 ${MAX_TRY}`
  do
  	filename=`get_pass_file`
  	if [ -f ${filename} ]; then
  		sleep 1
  		continue
  	else
  		findname=1
  		break
  	fi
  done

  if [ ${findname} -eq 0 ]; then
	return 1
  fi

  if ! touch ${filename}; then
	return 1
  fi

  echo "${filename}"
  return 0

}

#
# Utils to make pass to passwd file
# Usage: $0 <passfile> <passwd>
#
function make_pass_to_file()
{

  if [ $# -ne 2 ]; then
    return 2
  fi

  local passfile=$1
  local password=$2

  # use vncpasswd
  echo ${password} | ${VNC_PASSWD} -f > ${passfile} 2>/dev/null
  if [ $? -ne 0 ]; then
  	echo ${password} | vncpasswd -f > ${passfile} 2>/dev/null
  	if [  $? -ne 0 ]; then
  		test -f ${passfile} && rm -f ${passfile} > /dev/null 2>&1
		return 1
  	fi
  fi
  
  # use vncpasswd.arg
  #${VNC_PASSWD} ${passfile} ${password} 2>/dev/null
  
  chmod 600 ${passfile}
  #echo "passfile:${passfile}"

  return 0
}

#get all Vnc Proxy Servers from database on tables resource and resource_relation;
getVncProxyServers(){
    RES=`ssh -o StrictHostKeyChecking=no -o BatchMode=yes ${VNC_HOSTNAME} "/opt/gvmysql/bin/mysql --protocol=tcp -P 3309 -u${DB_USERNAME} -p${DB_PASSWD}"  <<EOF
        use ${DB_NAME};
        SELECT NAME FROM gv_rm_resource res,gv_rm_resource_relation rel WHERE res.ID=rel.DEPENDANT_ID AND rel.DEPENDED_ID='1006';
EOF`
    if [ -z "$RES" ];then
#       echo "$RES is empty"
        RES=`hostname`
    fi  
    
    if [ -n "$RES" ];then
        READL= echo $RES | cut -d' ' -f 2-  
    fi  
    echo $READL
}

